import models
from torch import optim,nn
from torch.nn import Module
from utils.parameters import *
from torchvision.models import resnet18

#the task template
class Task:
    params:Params = None
    train_dataset = None
    train_dataloader = None
    test_dataset = None
    test_dataloader = None
    model:Module = None
    optimizer:optim.Optimizer = None
    loss_fn = None

    def __init__(self, params:Params) -> None:
        self.params = params
        self._init_task()

    def _init_task(self) -> None:
        self.load_data()
        self.build_model(self.params.dst,self.params.model)
        self.model = self.model.to(self.params.device)
        self.optimizer = self.make_optimizer()
        self.loss_fn = self.make_loss_fn()

    def load_data(self) -> None:
        raise NotImplemented

    def make_optimizer(self, model=None) -> optim.Optimizer:
        if model is None:
            model = self.model
        if self.params.optimizer.lower() == 'sgd':
            optimizer = optim.SGD(model.parameters(),
                                  lr=self.params.lr,
                                  weight_decay=self.params.decay,
                                  momentum=self.params.momentum)
        elif self.params.optimizer.lower() == 'adam':
            optimizer = optim.Adam(model.parameters(),
                                   lr=self.params.lr,
                                   weight_decay=self.params.decay)
        else:
            raise ValueError(f'No optimizer: {self.params.optimizer}')
        return optimizer


    def make_loss_fn(self) -> Module:
        return nn.CrossEntropyLoss().to(self.params.device)
    
    def build_model(self,dst:str,model:str):
        if model.lower() == "lenet":
            self.model = models.LeNet(dst)
        elif model.lower() == "cnn4":
            self.model = models.CNN4(dst)
        elif model.lower() == "cnn2":
            self.model = models.CNN2(dst)
        elif model.lower() == "resnet18":
            self.handle_resnet_18(
                dst = dst,
                pretrained = self.params.pretrained,
                freeze_ftrs = self.params.freeze_ftrs
                )
        else:
            assert False, "模型名称错误"

    def handle_resnet_18(self,dst:str,num_classes:int = 20,pretrained:bool = False,freeze_ftrs:bool = False):
        _num_classes:int = num_classes
        if dst == "cifar10":
            _num_classes = 10
        self.model = resnet18(num_classes = _num_classes,pretrained=pretrained)

        if freeze_ftrs:#是否冻结除全连接层之外的参数
            for param in self.model.parameters():
                param.requires_grad = False
            num_ftrs = self.model.fc.in_features
            self.model.fc = nn.Sequential(nn.Linear(num_ftrs, num_classes))